// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/prestosql/aggregates/sum_aggregate.h>
#include <pollux/functions/lib/aggregates/sum_aggregate_base.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>

using namespace kumo::pollux::functions::aggregate;

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        template<typename TInput, typename TAccumulator, typename ResultType>
        using SumAggregate = SumAggregateBase<TInput, TAccumulator, ResultType, false>;

        template<template <typename U, typename V, typename W> class T>
        exec::AggregateRegistrationResult registerSum(
            const std::string &name,
            bool withCompanionFunctions,
            bool overwrite) {
            std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures{
                exec::AggregateFunctionSignatureBuilder()
                .returnType("real")
                .intermediateType("double")
                .argumentType("real")
                .build(),
                exec::AggregateFunctionSignatureBuilder()
                .returnType("double")
                .intermediateType("double")
                .argumentType("double")
                .build(),
                exec::AggregateFunctionSignatureBuilder()
                .integerVariable("a_precision")
                .integerVariable("a_scale")
                .argumentType("DECIMAL(a_precision, a_scale)")
                .intermediateType("VARBINARY")
                .returnType("DECIMAL(38, a_scale)")
                .build(),
            };

            for (const auto &inputType:
                 {
                     "tinyint",
                     "smallint",
                     "integer",
                     "bigint",
                     INTERVAL_DAY_TIME()->name()
                 }) {
                signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                    .returnType("bigint")
                    .intermediateType("bigint")
                    .argumentType(inputType)
                    .build());
            }

            return exec::registerAggregateFunction(
                name,
                std::move(signatures),
                [name](
            core::AggregationNode::Step step,
            const std::vector<TypePtr> &argTypes,
            const TypePtr &resultType,
            const core::QueryConfig & /*config*/)
            -> std::unique_ptr<exec::Aggregate> {
                    POLLUX_CHECK_EQ(argTypes.size(), 1, "{} takes only one argument", name);
                    auto inputType = argTypes[0];
                    switch (inputType->kind()) {
                        case TypeKind::TINYINT:
                            return std::make_unique<T<int8_t, int64_t, int64_t> >(BIGINT());
                        case TypeKind::SMALLINT:
                            return std::make_unique<T<int16_t, int64_t, int64_t> >(BIGINT());
                        case TypeKind::INTEGER:
                            return std::make_unique<T<int32_t, int64_t, int64_t> >(BIGINT());
                        case TypeKind::BIGINT: {
                            if (inputType->isShortDecimal()) {
                                return std::make_unique<DecimalSumAggregate<int64_t> >(resultType);
                            }
                            return std::make_unique<T<int64_t, int64_t, int64_t> >(BIGINT());
                        }
                        case TypeKind::HUGEINT: {
                            if (inputType->isLongDecimal()) {
                                return std::make_unique<DecimalSumAggregate<int128_t> >(
                                    resultType);
                            }
                            POLLUX_NYI();
                        }
                        case TypeKind::REAL:
                            if (resultType->kind() == TypeKind::REAL) {
                                return std::make_unique<T<float, double, float> >(resultType);
                            }
                            return std::make_unique<T<float, double, double> >(DOUBLE());
                        case TypeKind::DOUBLE:
                            if (resultType->kind() == TypeKind::REAL) {
                                return std::make_unique<T<double, double, float> >(resultType);
                            }
                            return std::make_unique<T<double, double, double> >(DOUBLE());
                        case TypeKind::VARBINARY:
                            // Always use int128_t template for Varbinary as the result
                            // type is either int128_t or
                            // UnscaledLongDecimalWithOverflowState.
                            return std::make_unique<DecimalSumAggregate<int128_t> >(resultType);
                        default:
                            POLLUX_UNREACHABLE(
                                "Unknown input type for {} aggregation {}",
                                name,
                                inputType->kindName());
                    }
                },
                {false /*orderSensitive*/, false /*companionFunction*/},
                withCompanionFunctions,
                overwrite);
        }
    } // namespace

    void registerSumAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        registerSum<SumAggregate>(prefix + kSum, withCompanionFunctions, overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
